///
/// @file CodeGeneratorArm64.cpp
/// @brief ARM64的后端处理实现
/// @author zenglj (zenglj@live.com)
/// @version 1.0
/// @date 2024-11-21
///
/// @copyright Copyright (c) 2024
///
/// @par 修改日志:
/// <table>
/// <tr><th>Date       <th>Version <th>Author  <th>Description
/// <tr><td>2024-11-21 <td>1.0     <td>zenglj  <td>新做
/// </table>
///
#include <cstdint>
#include <cstdio>
#include <string>
#include <vector>

#include "Function.h"
#include "Module.h"
#include "PlatformArm64.h"
#include "CodeGeneratorArm64.h"
#include "InstSelectorArm64.h"
#include "SimpleRegisterAllocator.h"
#include "ILocArm64.h"
#include "RegVariable.h"
#include "FuncCallInstruction.h"
#include "ArgInstruction.h"
#include "MoveInstruction.h"

/// @brief 构造函数
/// @param tab 符号表
CodeGeneratorArm64::CodeGeneratorArm64(Module * _module) : CodeGeneratorAsm(_module)
{}

/// @brief 析构函数
CodeGeneratorArm64::~CodeGeneratorArm64()
{}

/// @brief 产生汇编头部分
void CodeGeneratorArm64::genHeader()
{
    fprintf(fp, "%s\n", ".arch armv8-a");
}

/// @brief 全局变量Section，主要包含初始化的和未初始化过的
void CodeGeneratorArm64::genDataSection()
{
    // 输出未初始化的全局变量到.bss段
    bool hasBss = false;
    for (auto var: module->getGlobalVariables()) {
        if (var->isInBSSSection()) {
            if (!hasBss) {
                fprintf(fp, ".bss\n");
                hasBss = true;
            }
            fprintf(fp, ".comm %s, %d, %d\n", var->getName().c_str(), var->getType()->getSize(), var->getAlignment());
        }
    }

    // 输出已初始化的全局变量到.data段
    bool hasData = false;
    for (auto var: module->getGlobalVariables()) {
        if (!var->isInBSSSection()) {
            if (!hasData) {
                fprintf(fp, ".data\n");
                hasData = true;
            }
            fprintf(fp, ".align %d\n", var->getAlignment());
            fprintf(fp, ".global %s\n", var->getName().c_str());
            fprintf(fp, ".type %s, %%object\n", var->getName().c_str());
            fprintf(fp, "%s:\n", var->getName().c_str());

            auto initializer = var->getInitializer();
            if (initializer) {
                if (auto arr = dynamic_cast<ConstantArray *>(initializer)) {
                    // 全局数组
                    auto * arrType = static_cast<const PointerType *>(initializer->getType());
                    const Type * elemType = arrType->getElementType();
                    for (const auto & v: arr->getValues()) {
                        if (elemType->isIntegerType() || elemType->isCharType()) {
                            auto * ci = dynamic_cast<ConstInt *>(v);
                            fprintf(fp, "    .quad %lld\n", ci ? (long long) ci->getVal() : 0);
                        } else if (elemType->isFloatType()) {
                            auto * cf = dynamic_cast<ConstFloat *>(v);
                            fprintf(fp, "    .double %f\n", cf ? cf->getFloatVal() : 0.0f);
                        }
                    }
                } else {
                    // 单个全局变量
                    if (var->getType()->isIntegerType() || var->getType()->isCharType()) {
                        fprintf(fp, "    .quad %d\n", var->getInitIntValue());
                    } else if (var->getType()->isFloatType()) {
                        auto cf = dynamic_cast<ConstFloat *>(initializer);
                        fprintf(fp, "    .double %f\n", cf ? cf->getFloatVal() : 0.0f);
                    }
                }
            } else {
                fprintf(fp, "\t.zero %d\n", var->getType()->getSize());
            }
        }
    }

    // 生成代码段
    fprintf(fp, ".text\n");
}

///
/// @brief 获取IR变量相关信息字符串
/// @param str
///
void CodeGeneratorArm64::getIRValueStr(Value * val, std::string & str)
{
    std::string name = val->getName();
    std::string IRName = val->getIRName();
    int32_t regId = val->getRegId();
    int32_t baseRegId;
    int64_t offset;
    std::string showName;

    if (name.empty() && (!IRName.empty())) {
        showName = IRName;
    } else if ((!name.empty()) && IRName.empty()) {
        showName = name;
    } else if ((!name.empty()) && (!IRName.empty())) {
        showName = name + ":" + IRName;
    } else {
        showName = "";
    }

    if (regId != -1) {
        // 寄存器
        str += "\t@ " + showName + ":" + PlatformArm64::regName[regId];
    } else if (val->getMemoryAddr(&baseRegId, &offset)) {
        // 栈内寻址，[fp,#4]
        str += "\t@ " + showName + ":[" + PlatformArm64::regName[baseRegId] + ",#" + std::to_string(offset) + "]";
    }
}

/// @brief 针对函数进行汇编指令生成，放到.text代码段中
/// @param func 要处理的函数
void CodeGeneratorArm64::genCodeSection(Function * func)
{
    // 寄存器分配以及栈内局部变量的站内地址重新分配
    registerAllocation(func);

    // 获取函数的指令列表
    std::vector<Instruction *> & IrInsts = func->getInterCode().getInsts();

    // 汇编指令输出前要确保Label的名字有效，必须是程序级别的唯一，而不是函数内的唯一。要全局编号。
    for (auto inst: IrInsts) {
        if (inst->getOp() == IRInstOperator::IRINST_OP_LABEL) {
            inst->setName(IR_LABEL_PREFIX + std::to_string(labelIndex++));
        }
    }

    // ILOC代码序列
    ILocArm64 iloc(module);

    // 指令选择生成汇编指令
    InstSelectorArm64 instSelector(IrInsts, iloc, func, simpleRegisterAllocator);
    instSelector.setShowLinearIR(this->showLinearIR);
    instSelector.run();

    // 删除无用的Label指令
    iloc.deleteUsedLabel();

    // ILOC代码输出为汇编代码
    fprintf(fp, ".align %d\n", func->getAlignment());
    fprintf(fp, ".global %s\n", func->getName().c_str());
    fprintf(fp, ".type %s, %%function\n", func->getName().c_str());
    fprintf(fp, "%s:\n", func->getName().c_str());

    // 开启时输出IR指令作为注释
    if (this->showLinearIR) {

        // 输出有关局部变量的注释，便于查找问题
        for (auto localVar: func->getVarValues()) {
            std::string str;
            getIRValueStr(localVar, str);
            if (!str.empty()) {
                fprintf(fp, "%s\n", str.c_str());
            }
        }

        // 输出指令关联的临时变量信息
        for (auto inst: func->getInterCode().getInsts()) {
            if (inst->hasResultValue()) {
                std::string str;
                getIRValueStr(inst, str);
                if (!str.empty()) {
                    fprintf(fp, "%s\n", str.c_str());
                }
            }
        }
    }

    iloc.outPut(fp);
}

/// @brief 寄存器分配
/// @param func 函数指针
void CodeGeneratorArm64::registerAllocation(Function * func)
{
    // 内置函数不需要处理
    if (func->isBuiltin()) {
        return;
    }

    // 最简单/朴素的寄存器分配策略：局部变量和临时变量都保存在栈内，全局变量在静态存储.data区中
    // R0,R1,R2和R3寄存器不需要保护，可直接使用
    // SP寄存器预留，不需要保护，但需要保证值的正确性
    // R4-R10, fp(11), lx(14)都需要保护，没有函数调用的函数可不用保护lx寄存器
    // 被保留的寄存器主要有：
    //  (1) FP寄存器用于栈寻址，即R11
    //  (2) LX寄存器用于函数调用，即R14。没有函数调用的函数可不用保护lx寄存器
    //  (3) R10寄存器用于立即数过大时要通过寄存器寻址，这里简化处理进行预留

    std::vector<int32_t> & protectedRegNo = func->getProtectedReg();
    protectedRegNo.clear();
    protectedRegNo.push_back(ARM64_FP_REG_NO);
    if (func->getExistFuncCall()) {
        protectedRegNo.push_back(ARM64_LR_REG_NO);
    }

    // 调整函数调用指令，主要是前四个寄存器传值，后面用栈传递
    // 为了更好的进行寄存器分配，可以进行对函数调用的指令进行预处理
    // 当然也可以不做处理，不过性能更差。这个处理是可选的。
    adjustFuncCallInsts(func);

    // 为局部变量和临时变量在栈内分配空间，指定偏移，进行栈空间的分配
    stackAlloc(func);

    // 函数形参要求前四个寄存器分配，后面的参数采用栈传递，实现实参的值传递给形参
    // 这一步是必须的
    adjustFormalParamInsts(func);
}

/// @brief 寄存器分配前对函数内的指令进行调整，以便方便寄存器分配
/// @param func 要处理的函数
void CodeGeneratorArm64::adjustFormalParamInsts(Function * func)
{
    // 请注意这里所得的所有形参都是对应的实参的值关联的临时变量
    // 如果不是不能使用这里的代码
    auto & params = func->getParams();

    // 形参的前八个通过寄存器来传值R0-R7
    for (int k = 0; k < (int) params.size() && k <= 7; k++) {

        // 前八个设置分配寄存器
        simpleRegisterAllocator.bitmapSet(k);
        params[k]->setRegId(k);
    }

    // 根据ARM版C语言的调用约定，除前8个外的实参进行值传递，逆序入栈
    int64_t fp_esp = func->getMaxDep() + (func->getProtectedReg().size() * 8);
    // 16字节对齐形参空间起始位置
    fp_esp += (16 - fp_esp % 16) % 16;
    
    for (int k = 8; k < (int) params.size(); k++) {
        // 8字节对齐每个形参
        int32_t param_size = params[k]->getType()->getSize();
        param_size += (8 - param_size % 8) % 8;

        params[k]->setMemoryAddr(ARM64_FP_REG_NO, fp_esp);

        // 增加对齐后的大小
        fp_esp += param_size;
    }
}

/// @brief 寄存器分配前对函数内的指令进行调整，以便方便寄存器分配
/// @param func 要处理的函数
void CodeGeneratorArm64::adjustFuncCallInsts(Function * func)
{
    auto & insts = func->getInterCode().getInsts();

    for (auto pIter = insts.begin(); pIter != insts.end(); ++pIter) {
        if (Instanceof(callInst, FuncCallInstruction *, *pIter)) {
            int paramNum = callInst->getOperandsNum();
            if (paramNum < 1)
                continue; // 至少有retVal

            int argCount = paramNum - 1; // 真正参数个数

            // 处理超过8个参数，逆序入栈（参数索引 0..argCount-1）
            // 16字节对齐实参空间
            int32_t stack_offset = 0;
            for (int32_t k = argCount - 1; k >= 8; k--) {
                auto arg = callInst->getOperand(k);
                if (!arg)
                    continue;

                // 16字节对齐每个实参
                int32_t arg_size = arg->getType()->getSize();
                arg_size += (8 - arg_size % 8) % 8;

                LocalVariable * newVal = func->newLocalVarValue(IntegerType::getTypeInt());
                newVal->setMemoryAddr(ARM64_SP_REG_NO, stack_offset);
                stack_offset += arg_size;

                Instruction * assignInst = new MoveInstruction(func, newVal, arg);
                callInst->setOperand(k, newVal);
                pIter = insts.insert(pIter, assignInst);
                ++pIter;
            }

            // 加上整体对齐
            stack_offset += (16 - stack_offset % 16) % 16;

            // 前8个参数用寄存器 (参数索引 0..7)
            for (int k = 0; k < argCount && k < 8; k++) {
                auto arg = callInst->getOperand(k);
                if (!arg)
                    continue;
                // 如果已经在对应寄存器就跳过
                if (arg->getRegId() == k)
                    continue;
                Instruction * assignInst = new MoveInstruction(func, PlatformArm64::intRegVal[k], arg);
                simpleRegisterAllocator.bitmapSet(k);
                callInst->setOperand(k, PlatformArm64::intRegVal[k]);
                pIter = insts.insert(pIter, assignInst);
                ++pIter;
            }

            // 插入ARG指令，参数索引 0..argCount-1
            for (int k = 0; k < argCount; k++) {
                auto arg = callInst->getOperand(k);
                if (!arg)
                    continue;
                pIter = insts.insert(pIter, new ArgInstruction(func, arg));
                ++pIter;
            }

            // 返回值处理，返回值变量是最后一个操作数
            if (callInst->hasResultValue()) {
                // 如果返回值变量的寄存器ID不是0（x0），则插入寄存器到变量的move指令
                if (callInst->getRegId() != 0) {
                    Instruction * assignInst = new MoveInstruction(func, callInst, PlatformArm64::intRegVal[0]);
                    pIter = insts.insert(pIter + 1, assignInst);
                }
            }
        }
    }
}

/// @brief 栈空间分配
/// @param func 要处理的函数
void CodeGeneratorArm64::stackAlloc(Function * func)
{
    int32_t sp_esp = 0;

    // 获取函数变量列表
    std::vector<LocalVariable *> & vars = func->getVarValues();

    for (auto var: vars) {
        if ((var->getRegId() == -1) && (!var->getMemoryAddr())) {
            int32_t size = var->getType()->getSize();
            size += (8 - size % 8) % 8; // 保持16字节对齐
            var->setMemoryAddr(ARM64_FP_REG_NO, sp_esp);
            sp_esp += size;
        }
    }

    for (auto inst: func->getInterCode().getInsts()) {
        if (inst->hasResultValue()) {
            int32_t size = inst->getType()->getSize();
            size += (8 - size % 8) % 8; // 保持16字节对齐
            inst->setMemoryAddr(ARM64_FP_REG_NO, sp_esp);
            sp_esp += size;
        }
    }

    sp_esp += (16 - sp_esp % 16) % 16; // 确保整个栈帧16字节对齐
    func->setMaxDep(sp_esp);
}